import numpy as np
import pyopencl as cl
from time import time


def run(length):
    print(__file__)
    print("length:", length)
    rng = np.random.default_rng()
    a_np = rng.random(length, dtype=np.float32)
    b_np = rng.random(length, dtype=np.float32)
    ctx = cl.create_some_context()
    queue = cl.CommandQueue(ctx)
    mf = cl.mem_flags
    a_g = cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=a_np)
    b_g = cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=b_np)

    prg = cl.Program(ctx, """
    __kernel void sum(
        __global const float2 *a_g, __global const float2 *b_g, __global float2 *res_g)
    {
      int gid = get_global_id(0);
      res_g[gid] = a_g[gid] + b_g[gid];
    }
    """).build()

    res_g = cl.Buffer(ctx, mf.WRITE_ONLY, a_np.nbytes)
    knl = prg.sum  # Use this Kernel object for repeated calls
    tcl1 = time()
    knl(queue, a_np.shape, None, a_g, b_g, res_g)
    tcl2 = time()

    # opencl while test
    # while True:
    #     knl(queue, a_np.shape, None, a_g, b_g, res_g)

    res_np = np.empty_like(a_np)
    cl.enqueue_copy(queue, res_np, res_g)

    # Check on CPU with Numpy:
    tnp1 = time()
    np_sum = a_np + b_np
    tnp2 = time()
    error_np = res_np - (a_np + b_np)
    print(f"Error:\n{error_np}")
    print(f"Norm: {np.linalg.norm(error_np):.16e}")
    assert np.allclose(res_np, a_np + b_np)

    # Check on CPU with native python
    res_na = np.empty_like(a_np)
    tna1 = time()
    for i in range(length):
        res_na[i] = a_np[i] + b_np[i]
    tna2 = time()

    # time
    print(f"Python : {tna2 - tna1:.10f}s")
    print(f"Numpy  : {tnp2 - tnp1:.10f}s")
    print(f"OpenCL : {tcl2 - tcl1:.10f}s")


if __name__ == "__main__":
    run(50)
    run(500)
    run(5000)
    run(50000)
    run(500000)
    run(5000000)
    run(50000000)
